from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import HTMLResponse
from fastapi.templating import Jinja2Templates
from fastapi import Request
import uvicorn
from pydantic import BaseModel
from server.agent.agent_lingdongai import agent_prompt
from server.intent_recog.intent_recog_inference import recog
from server.rag import rag
from server.ner.ner_detection import ner
from server import neo4jserver
import json

app = FastAPI()

# 启用 CORS 中间件
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 可以指定具体的前端URL地址
    allow_credentials=True,
    allow_methods=["*"],  # 允许所有HTTP方法 (GET, POST, OPTIONS等)
    allow_headers=["*"],  # 允许所有请求头
)

# 初始化 Jinja2 模板
templates = Jinja2Templates(directory="templates")

# 自定义输入参数
class Item(BaseModel):
    info: str

# 聊天
@app.post("/task/chat")
def create_item(item: Item):

    # 给 response 赋予默认值
    response = None

    ## 意图识别
    # 如果是闲聊, 直接和大模型对话
    single_intention = recog(item.info)
    if int(single_intention) == 0:
        print("进入闲聊")
        response = agent_prompt.qwen(item.info, "")
    # 如果是问诊, 调用rag+知识图谱
    else:
        print("进入问诊")
        # 调用rag
        chat = rag.RAG()
        rag_info = chat.askAndFindFiles(item.info)
        print("rag检索到的信息: ", rag_info)
        # 调用命名实体识别
        ner_info = ner(item.info)
        print("医学命名实体识别出的结果: ", ner_info)
        # 调用知识图谱
        # neo4j_info = neo4jserver.fetch_person(ner_info)
        # print("知识图谱查询结果: ", neo4j_info)

        # 将rag检索信息和知识图谱结合
        prompt = f"{rag_info} \n 并根据疾病给出对应的科室"
        print("提示词工程: ", prompt)
        
        response = agent_prompt.qwen(item.info, prompt)

    return {"message": f"{response}"}

# 打开 HTML 页面接口
@app.get("/", response_class=HTMLResponse)
def read_html(request: Request):
    return templates.TemplateResponse("index.html", {"request": request})

# 如果通过脚本启动，则使用 uvicorn.run 指定接口和端口
if __name__ == "__main__":
    print("全部程序启动完毕")
    uvicorn.run(app, host="0.0.0.0", port=8000)
